import torch
from pyDOE import lhs
from torch import Tensor

from config import device
from individual.population import Population
from utils.annotation import tensor_to_2dim


class Problem:
    def __init__(self, var_dim, obj_dim, low_limit, high_limit, max_fun_eval, kwargs):
        self.var_dim = var_dim
        self.obj_dim = obj_dim
        self.low_limit = low_limit
        self.high_limit = high_limit
        self.max_fun_eval = max_fun_eval
        self.fun_eval = 0
        self.parse_params(kwargs)
        self.interval_length = None
        self.set_boundary(low_limit, high_limit)
        self.optimal_solutions = self.get_optimal_solutions(10000)

    def set_boundary(self, low_limit, high_limit):
        if low_limit is not None and high_limit is not None:
            self.low_limit = low_limit
            self.high_limit = high_limit
            self.interval_length = high_limit - low_limit

    def parse_params(self, kwargs):
        """
        将参数设置到self上
        :param kwargs:
        :return:
        """
        for k, v in kwargs.items():
            self.__setattr__(k, v)

    @tensor_to_2dim
    def repair_decision(self, dec):
        return dec.clamp(self.low_limit, self.high_limit)

    def eval_value(self, x: Tensor):
        pass

    def estimate_population(self, dec, **kwargs):
        """
        评估种群
        :param dec:
        :param kwargs:
        :return:
        """
        dec = self.repair_decision(dec)
        index = torch.isnan(dec)
        if index.any():
            dec[index] = (torch.rand(dec.shape, dtype=dec.dtype) * self.interval_length)[index]
        self.fun_eval += len(dec)
        obj = self.eval_value(dec)
        # index = unique(obj)
        # obj = obj[index]
        # dec = dec[index]
        return Population(dec, obj, con=self.cal_con(dec), **kwargs)

    def cal_con(self, dec):
        return torch.zeros(len(dec), 1, dtype=torch.double)

    def __eq__(self, other):
        for k, v in self.__dict__.items():
            if other.__dict__[k] != v:
                return False
        return True

    def get_optimal_solutions(self, size):
        pass

    def init_population(self, pop_size, **kwargs):
        dec = torch.tensor(lhs(samples=pop_size, n=self.var_dim), dtype=torch.double, device=device) * (
                self.high_limit - self.low_limit) + self.low_limit
        return self.estimate_population(dec, **kwargs)

    def cal_grad(self, dec):
        """
        对单个个体进行求解近似导数
        :param dec: 1*var_dim
        :return:
        """
        dec[dec == 0] = 1e-12
        x = dec * (1 + torch.eye(len(dec)) * 1e-6)
        return (self.eval_value(x) - self.estimate_population(dec)[0].obj) / dec.unsqueeze(dim=1) / 1e-6

    @tensor_to_2dim
    def calc_grad(self, dec):
        """
        对种群进行求解精确导数
        :param dec:
        :return: 返回三维导数张量，大小为：dec.size(0)*dec.size(1)*obj_dim
        """
        grad_list = []
        dec.requires_grad_()
        obj = self.eval_value(dec)
        # 创建一个梯度张量，形状与输出张量 y 相同
        grad_outputs = torch.ones_like(obj)
        # 计算偏导
        for dim in range(self.obj_dim):
            grad = torch.autograd.grad(outputs=obj[:, dim], inputs=dec,
                                       grad_outputs=grad_outputs[:, dim], retain_graph=True)[0]
            grad_list.append(grad)
        return torch.stack(grad_list, dim=2)
